# 题目一：使用基于混合高斯模型的背景提取算法，提取前景并显示(显示二值化图像，前景为白色)
import cv2

videoFileName = 'data/vtest.avi'

cap = cv2.VideoCapture(videoFileName)  # 打开视频，采集到帧
fgbg = cv2.createBackgroundSubtractorMOG2()  # 该函数用于背景提取
thresh = 200

while True:  
    ret, frame = cap.read()  # 提取视频中的每一帧
    if not ret:  # 没读到当前帧，结束
        break
    
    fgmask = fgbg.apply(frame)  # apply函数：对当前帧进行背景提取，并计算出前景对应的mask图像
    #  由于前景存在一些噪声，用threshold函数把当前帧的前景图像进行初步过滤    
    _, fgmask = cv2.threshold(fgmask, 30, 0xff, cv2.THRESH_BINARY)  
    bgImage = fgbg.getBackgroundImage()  # 得到实时更新后的背景图像
    #  首先调用findContours对当前的图像进行分割，得到的分割结果（目标）放到变量cnts里   
    #  注:cv2.findContours()函数接受的参数为二值图，即黑白的（不是灰度图），所以读取的图像要先转成灰度的，再转成二值图
    binary , contours, hierarchy = cv2.findContours(fgmask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # cv2.RETR_EXTERNAL只检测外轮廓

    count = 0
    colors = [(0, 0xff, 0), (255,182,193), (106,90,205), (220,20,60), (0,255,255), (199,21,133), (147,112,219), (0,0,255), (100,149,237), (255,255,0)]
    for c in contours:  # 处理每个目标
        area = cv2.contourArea(c)
        if area < thresh:  # 区域面积小于指定阈值，认为它是噪声
            continue
        count += 1
        
        # 题目二：在1基础上，将前景目标进行分割，进一步使用不同颜色矩形框标记，并在命令行窗口中输出每个矩形框的位置和大小。 
        x, y, w, h = cv2.boundingRect(c)
        cv2.rectangle(frame, (x,y), (x+w, y+h), colors[count-1], 2)  # 把检测到的目标通过矩形块画出来
        cv2.rectangle(fgmask, (x,y), (x+w, y+h), colors[count-1], 2)  # 把检测到的目标通过矩形块画出来
        print('第', count, '个目标的矩形框位置为：',(x, y), '大小为：', w, '*', h)

    print('共检测到', count, '个目标', '\n')
    cv2.imshow('frame', frame)  # 当前帧
    cv2.imshow('Background', bgImage)  # 背景
    cv2.imshow('object', fgmask)  # 前景
    
    key = cv2.waitKey(30)  # 每一帧间隔30ms
    if key == 27:  # 按下 ESC 键，退出
        break

cap.release()
cv2.destroyAllWindows()



# 扩展作业（作业4）：使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。 
videoFileName = 'data/vtest.avi'
# 角点检测参数
feature_params = dict( maxCorners = 100,
                       qualityLevel = 0.3,
                       minDistance = 7,
                       blockSize = 7)
# lucas kanade 光流法参数
lk_params = dict( winSize = (15, 15),
                  maxLevel = 2,
                  criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
cap = cv2.VideoCapture(videoFileName)

# 计算第一帧特征点
ret, prev = cap.read()  # 读取第一帧
prevGray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
p0 = cv2.goodFeaturesToTrack(prevGray, mask=None, **feature_params)  # 得到较好的特征点

while True:
    ret, frame = cap.read()  # 取每一帧
    if not ret:  # 没读到当前帧，结束
        break
    
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)  # 转灰度图像
    
    # 计算光流
    p1, st, err = cv2.calcOpticalFlowPyrLK(prevGray, gray, p0, None, **lk_params)  # p1为特征点，st为每个特征点的置信度
    # 选取好的跟踪点
    goodPoints = p1[st==1]
    goodPrevPoints = p0[st==1]
    
    # 在结果图像中跌加画出特征点和计算出来的光流向量
    res = frame.copy()
    drawColor = (0, 0, 255)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv2.line(res, (x0, y0), (x1, y1), drawColor)
        cv2.circle(res, (x0, y0), 3, drawColor)
    
    # 更新上一帧
    prevGray = gray.copy()
    p0 = goodPoints.reshape(-1, 1, 2)
    
    # 显示计算结果图像
    cv2.imshow('res', res)
    
    key = cv2.waitKey(30)
    if key == 27:
        break